package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

// 日志文件路径
const logFile = "logs/agent.log"

// API响应结构
type Response struct {
	Code    int         `json:"code"`
	Message string      `json:"message"`
	Data    interface{} `json:"data,omitempty"`
}

// 更新配置请求结构
type UpdateConfigRequest struct {
	ConfigContent string `json:"config"`
	ReloadURL     string `json:"reloadURL"`
	Username      string `json:"username"`
	Password      string `json:"password"`
	ConfigPath    string `json:"configPath"`
}

// min函数
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

// 创建备份目录
func createBackupDir(configPath string) (string, error) {
	configDir := filepath.Dir(configPath)
	backupDir := filepath.Join(configDir, "backups")
	if err := os.MkdirAll(backupDir, 0755); err != nil {
		return "", fmt.Errorf("创建备份目录失败: %v", err)
	}
	return backupDir, nil
}

// 配置日志
func setupLogger() *os.File {
	logDir := filepath.Dir(logFile)
	if err := os.MkdirAll(logDir, 0755); err != nil {
		log.Fatalf("创建日志目录失败: %v", err)
	}

	file, err := os.OpenFile(logFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatalf("打开日志文件失败: %v", err)
	}

	log.SetOutput(file)
	log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

	return file
}

// 检查Prometheus服务可用性
func checkPrometheusService(baseURL string) error {
	if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") {
		baseURL = "http://" + baseURL
	}

	u, err := url.Parse(baseURL)
	if err != nil {
		return fmt.Errorf("URL格式无效: %v", err)
	}

	hostPort := u.Host
	if !strings.Contains(hostPort, ":") {
		hostPort = hostPort + ":9090"
	}

	conn, err := net.DialTimeout("tcp", hostPort, 5*time.Second)
	if err != nil {
		return fmt.Errorf("端口 %s 不可访问: %v", hostPort, err)
	}
	conn.Close()
	return nil
}

// 重新加载Prometheus配置
func reloadPrometheus(reloadURL, username, password string) error {
	if reloadURL == "" {
		return fmt.Errorf("重载URL不能为空")
	}

	if !strings.HasPrefix(reloadURL, "http://") && !strings.HasPrefix(reloadURL, "https://") {
		reloadURL = "http://" + reloadURL
	}

	if !strings.HasSuffix(reloadURL, "/-/reload") {
		if strings.HasSuffix(reloadURL, "/") {
			reloadURL = reloadURL + "-/reload"
		} else {
			reloadURL = reloadURL + "/-/reload"
		}
	}

	if err := checkPrometheusService(reloadURL); err != nil {
		return fmt.Errorf("Prometheus服务检查失败: %v", err)
	}

	client := &http.Client{Timeout: 30 * time.Second}
	req, err := http.NewRequest("POST", reloadURL, nil)
	if err != nil {
		return fmt.Errorf("创建请求失败: %v", err)
	}

	if username != "" && password != "" {
		req.SetBasicAuth(username, password)
	}

	resp, err := client.Do(req)
	if err != nil {
		return fmt.Errorf("发送重载请求失败: %v", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("重载请求失败,状态码 %d", resp.StatusCode)
	}

	return nil
}

func main() {
	port := flag.String("port", "8080", "agent监听端口")
	token := flag.String("token", "", "agent认证token，所有接口需携带Authorization: Bearer <token>")
	flag.Parse()

	agentToken := *token

	logFile := setupLogger()
	defer logFile.Close()

	r := gin.Default()

	// Bearer Token认证中间件
	r.Use(func(c *gin.Context) {
		reqToken := c.GetHeader("Authorization")
		expected := "Bearer " + agentToken
		if agentToken == "" || reqToken != expected {
			c.AbortWithStatusJSON(401, Response{Code: 401, Message: "Unauthorized: Bearer token required"})
			return
		}
		c.Next()
	})

	// 获取配置文件接口 - 用于统计功能
	r.GET("/get-config", func(c *gin.Context) {
		// 默认配置文件路径，可以通过环境变量或查询参数指定
		configPath := c.Query("configPath")
		if configPath == "" {
			configPath = "/etc/prometheus/prometheus.yml" // 默认路径
		}

		log.Printf("Reading config file: %s", configPath)
		content, err := ioutil.ReadFile(configPath)
		if err != nil {
			log.Printf("Error reading config file: %v", err)
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("读取配置文件失败: %v", err)})
			return
		}

		log.Printf("Config file content length: %d", len(content))
		log.Printf("Config file content preview: %s", string(content[:min(200, len(content))]))

		// 直接返回配置文件内容
		c.String(200, string(content))
	})

	// 读取配置文件接口
	r.POST("/config", func(c *gin.Context) {
		var req struct {
			ConfigPath string `json:"configPath" binding:"required"`
		}

		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(400, Response{Code: 400, Message: "无效的请求参数"})
			return
		}

		content, err := ioutil.ReadFile(req.ConfigPath)
		if err != nil {
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("读取配置文件失败: %v", err)})
			return
		}

		c.JSON(200, Response{Code: 200, Message: "成功", Data: string(content)})
	})

	// 更新配置文件接口
	r.POST("/update-config", func(c *gin.Context) {
		var req UpdateConfigRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(400, Response{Code: 400, Message: "无效的请求参数"})
			return
		}

		if req.ConfigPath == "" || req.ReloadURL == "" {
			c.JSON(400, Response{Code: 400, Message: "缺少必要参数: configPath, reloadUrl"})
			return
		}

		backupDir, err := createBackupDir(req.ConfigPath)
		if err != nil {
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("创建备份目录失败: %v", err)})
			return
		}
		timestamp := time.Now().Format("20060102150405")
		backupPath := filepath.Join(backupDir, fmt.Sprintf("prometheus.yml.bak.%s", timestamp))

		// 先备份原文件
		if err := copyFile(req.ConfigPath, backupPath); err != nil {
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("创建备份失败: %v", err)})
			return
		}

		// 写入新配置
		if err := ioutil.WriteFile(req.ConfigPath, []byte(req.ConfigContent), 0644); err != nil {
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("写入配置失败: %v", err)})
			return
		}

		// 重新加载Prometheus
		if err := reloadPrometheus(req.ReloadURL, req.Username, req.Password); err != nil {
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("重载Prometheus失败: %v", err)})
			return
		}
		c.JSON(200, Response{Code: 200, Message: "配置更新成功"})
	})

	// 新增：流式下载配置文件接口
	r.POST("/config/download", func(c *gin.Context) {
		var req struct {
			ConfigPath string `json:"configPath" binding:"required"`
		}
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(400, Response{Code: 400, Message: "无效的请求参数"})
			return
		}
		file, err := os.Open(req.ConfigPath)
		if err != nil {
			c.JSON(500, Response{Code: 500, Message: fmt.Sprintf("读取配置文件失败: %v", err)})
			return
		}
		defer file.Close()
		c.Header("Content-Type", "application/octet-stream")
		c.Header("Content-Disposition", "attachment; filename=prometheus.yml")
		c.Status(200)
		buf := make([]byte, 4096)
		for {
			n, err := file.Read(buf)
			if n > 0 {
				c.Writer.Write(buf[:n])
				c.Writer.Flush()
			}
			if err != nil {
				break
			}
		}
	})

	addr := fmt.Sprintf(":%s", *port)
	log.Printf("启动Prometheus agent，监听端口 %s", addr)
	if err := r.Run(addr); err != nil {
		log.Fatalf("启动服务器失败: %v", err)
	}
}

// 工具函数：文件拷贝
func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}
